#!/usr/bin/python

from random import randint
from unittest import TestCase as test_case, main
from cStringIO import StringIO

import sys
sys.path.extend(('../src/tools', '../src'))

from tools import generate_weighted_edges, load_graph

from heap import heap
from graph import find_components
from dtree import build_node, inner_node, leaf_node

from graph import Graph, find_path_prio

def create_2_components():
    edges = {
        (1, 2) : 1, (2, 1) : 1,
        (2, 3) : 1, (3, 2) : 1,
        (1, 3) : 1, (3, 1) : 1,
        (4, 5) : 1, (5, 4) : 1,
        (5, 6) : 1, (6, 5) : 1,
        (4, 6) : 1, (6, 4) : 1
        }
    return edges


class heap_test(test_case):

    @staticmethod
    def priority(x):
        return x

    def test_sample(self):
        from heapq import heappush, heappop

        mh = heap(self.priority)
        ph = []

        for iter in range(randint(100, 1100)):
            item = randint(1, 2000)

            heappush(ph, item)
            mh.push(item)

            self.assertTrue(ph[0] == mh.top())

        top = mh.top()
        while top is not None:
            self.assertTrue(top == heappop(ph))
            mh.pop()
            top = mh.top()

        self.assertTrue(ph == [])

    def test_empty(self):
        hp = heap(self.priority)

        for i in range(10):
            hp.push(randint(1, 100))

        for i in range(10):
            self.assertTrue(hp.top() is not None)
            hp.pop()

        for i in range(10):
            self.assertTrue(hp.top() is None)
            hp.pop()


class test_wgraph(test_case):

    def test_romania(self):
        p = ['Arad', 'Sibiu', 'Riminicu_Vilcea', 'Pitesti', 'Bucharest']
        with open('../data/romania.g') as fd:
            G = Graph(load_graph(fd))
            #G.load(fd)
            cost, path = find_path_prio(G, 'Arad', 'Bucharest')
            self.assertTrue(cost == 418)
            self.assertTrue(path == p)

    def test_sample(self):
        edges, path = generate_weighted_edges(2000, 1000)
        G = Graph(edges)

        cost, found_path = find_path_prio(G, path[0], path[-1])
        self.assertTrue(path == found_path)

    def test_components(self):
        G = Graph(create_2_components())
        cost, path = find_path_prio(G, 1, 5)
        self.assertTrue(cost is None)
        self.assertTrue(path is None)

    def test_splitting(self):
        G = Graph(create_2_components())
        g1, g2 = find_components(G)

        g10 = Graph({
            (1, 2) : 1, (2, 1) : 1,
            (2, 3) : 1, (3, 2) : 1,
            (1, 3) : 1, (3, 1) : 1})
        self.assertTrue(g1 == g10)

        g20 = Graph({
            (4, 5) : 1, (5, 4) : 1,
            (5, 6) : 1, (6, 5) : 1,
            (4, 6) : 1, (6, 4) : 1})
        self.assertTrue(g2 == g20)


class dtree_test(test_case):
    '''The data for this has been taken from the page

        http://decisiontrees.net/decision-trees-tutorial/tutorial-5-exercise-2/

    '''

    datastr = '''
# District House Type    Income  Previous   Outcome
#                                Customer
Suburban   Detached       High       No     Nothing
Suburban   Detached       High       Yes    Nothing
Rural      Detached       High       No     Responded
Urban      Semi-detached  High       No     Responded
Urban      Semi-detached  Low        No     Responded
Urban      Semi-detached  Low        Yes    Nothing
Rural      Semi-detached  Low        Yes    Responded
Suburban   Terrace        High       No     Nothing
Suburban   Semi-detached  Low        No     Responded
Urban      Terrace        Low        No     Responded
Suburban   Terrace        Low        Yes    Responded
Rural      Terrace        High       Yes    Responded
Rural      Detached       Low        No     Responded
Urban      Terrace        High       Yes    Nothing
'''

    def _read_nodes(self):
        L = [x.strip().split() for x in StringIO(self.datastr).readlines() if x and x[0] != '#']
        L = [(x[:-1], x[-1]) for x in L if x]
        attrset = set(range(len(L[0][0])))

        return L, attrset

    def test_buildID3(self):
        records, attrset = self._read_nodes()
        root = build_node(records, attrset)

        # check the root model
        self.assertTrue(isinstance(root, inner_node))
        self.assertTrue(root._attr == 0)

        rural = root._children['Rural']
        self.assertTrue(isinstance(rural, leaf_node))
        self.assertTrue(rural.value == 'Responded')

        suburban = root._children['Suburban']
        self.assertTrue(isinstance(suburban, inner_node))
        self.assertTrue(suburban._attr == 2)

        urban = root._children['Urban']
        self.assertTrue(isinstance(urban, inner_node))

        high = suburban._children['High']
        self.assertTrue(isinstance(high, leaf_node))
        self.assertTrue(high.value == 'Nothing')

        low = suburban._children['Low']
        self.assertTrue(isinstance(low, leaf_node))
        self.assertTrue(low.value == 'Responded')

        customer_no = urban._children['No']
        self.assertTrue(isinstance(customer_no, leaf_node))
        self.assertTrue(customer_no.value == 'Responded')

        customer_yes = urban._children['Yes']
        self.assertTrue(isinstance(customer_yes, leaf_node))
        self.assertTrue(customer_yes.value == 'Nothing')


if __name__ == '__main__':
    main()
